import mysql from 'mysql2/promise';
import { QdrantClient } from '@qdrant/js-client-rest';
import { log } from './logger';

// MySQL连接配置
const dbConfig = {
  host: process.env.MYSQL_HOST || 'localhost',
  port: parseInt(process.env.MYSQL_PORT || '3306'),
  user: process.env.MYSQL_USER || 'root',
  password: process.env.MYSQL_PASSWORD || '',
  database: process.env.MYSQL_DATABASE || 'alibaba_rag',
  waitForConnections: true,
  connectionLimit: 10,
  queueLimit: 0,
};

// 创建MySQL连接池
export const mysqlPool = mysql.createPool(dbConfig);

// Qdrant客户端配置
const qdrantConfig = {
  url: process.env.QDRANT_URL || 'http://localhost:6333',
  apiKey: process.env.QDRANT_API_KEY,
};

// 创建Qdrant客户端
export const qdrantClient = new QdrantClient(qdrantConfig);

// 数据库连接测试函数
export async function testMySQLConnection() {
  try {
    const connection = await mysqlPool.getConnection();
    await connection.ping();
    connection.release();
    log.info('MySQL连接成功');
    return true;
  } catch (error) {
    log.error('MySQL连接失败', error);
    return false;
  }
}

export async function testQdrantConnection() {
  try {
    const info = await qdrantClient.getCollections();
    log.info('Qdrant连接成功');
    return true;
  } catch (error) {
    log.error('Qdrant连接失败', error);
    return false;
  }
}

// 初始化Qdrant集合
export async function initializeQdrantCollection() {
  const collectionName = process.env.QDRANT_COLLECTION_NAME || 'documents';
  const embeddingDimensions = parseInt(
    process.env.EMBEDDING_DIMENSIONS || '1536'
  );

  try {
    // 检查集合是否存在
    const collections = await qdrantClient.getCollections();
    const collectionExists = collections.collections.some(
      collection => collection.name === collectionName
    );

    if (!collectionExists) {
      // 创建新集合
      await qdrantClient.createCollection(collectionName, {
        vectors: {
          size: embeddingDimensions,
          distance: 'Cosine',
        },
      });
      log.info(`Qdrant集合 '${collectionName}' 创建成功`);
    } else {
      log.info(`Qdrant集合 '${collectionName}' 已存在`);
    }

    return true;
  } catch (error) {
    log.error('Qdrant集合初始化失败', error);
    return false;
  }
}

// 数据库初始化函数
export async function initializeDatabases() {
  log.info('正在初始化数据库连接...');

  const mysqlOk = await testMySQLConnection();
  const qdrantOk = await testQdrantConnection();

  if (mysqlOk && qdrantOk) {
    await initializeQdrantCollection();
    log.info('所有数据库连接初始化完成');
    return true;
  } else {
    log.error('数据库连接初始化失败');
    return false;
  }
}
